import xml.etree.ElementTree as ET
import numpy as np
import glob
import os
import re

from time import time, localtime, strftime
from itertools import product

def configure_xml(xml_tree, model_path, texture_path, scale):
	root = xml_tree.getroot()
	texture_node = root.find('texture')
	# set our texture path
	texture_node.find('string').attrib["value"] = texture_path

	model_node = [node for node in root.iter('shape') if node.attrib['id'] == 'model_test'][0]
	model_node.find('string').attrib["value"] = model_path

	mat = model_node.find('transform').find('matrix').attrib['value']
	mat = np.array([float(v) for v in mat.split(' ')])

	mat[2] = -scale
	mat[4] = -scale
	mat[9] = scale

	mat_str = np.array2string(mat, max_line_width=1e9, precision=6, separator=' ', floatmode='fixed')
	mat_str = mat_str.replace('[ ', '').replace(']', '')
	model_node.find('transform').find('matrix').attrib['value'] = mat_str
	return xml_tree


def light_noise(xml_tree):
	root = xml_tree.getroot()

	light_nodes = [node for node in root.iter('shape') if node.attrib['id'].startswith("light")]

	for light_node in light_nodes:
		light_type = light_node.attrib["id"]
		mat = light_node.find('transform').find('matrix').attrib['value']
		mat = np.array([float(v) for v in mat.split(' ')])

		if "middle2" in light_type:
			mat[3] += np.random.random(1) - 0.5
			mat[7] += 0.2 * np.random.random(1) - 0.1
		else:
			mat[3] += 2 * np.random.random(1) - 1
			mat[7] += 2 * np.random.random(1) - 1

		mat[11] += 2 * np.random.random(1) - 1

		mat_str = np.array2string(mat, max_line_width=1e9, precision=6, separator=' ', floatmode='fixed')
		mat_str = mat_str.replace('[ ', '').replace(']', '')
		light_node.find('transform').find('matrix').attrib['value'] = mat_str

	return xml_tree


def generate_xml(resource_path, output_path, random_num=3):
	light_xmls = glob.glob(resource_path + "light_kuro/*.xml")
	models_path = glob.glob(resource_path + "model/*.serialized")
	# models_path = [resource_path + "model/db0.serialized"]
	textures_path = glob.glob(resource_path + "texture/stripe/*.jpg")

	if not os.path.exists(output_path):
		os.makedirs(output_path)

	light_and_texture = list(product(light_xmls, textures_path))

	for model_path in models_path:
		count = 0
		model_name = re.match(r'.*model\\(.*).serialized', model_path)[1]
		print("Processing ", model_name)

		xml_output = output_path + "xml/{}_stripe/".format(model_name)

		if not os.path.exists(xml_output):
			os.makedirs(xml_output)

		for light_xml, texture_path in light_and_texture:
			for i in range(random_num):
				xml_tree = ET.parse(light_xml)

				xml_tree = configure_xml(xml_tree, model_path, texture_path, 4)
				xml_tree = light_noise(xml_tree)
				
				xml_tree.write(xml_output + "{}.xml".format(count))
				count += 1

				if count % 100 == 0:
					print("Processed progress: ", count)

	print("Finished")


def generate_xml_sample(resource_path, output_path, sample_num):
	light_xmls = glob.glob(resource_path + "light_kuro/*.xml")
	models_path = glob.glob(resource_path + "model/*.serialized")
	# models_path = [resource_path + "model/db0.serialized"]
	textures_path = glob.glob(resource_path + "texture/**/*.jpg")

	if not os.path.exists(output_path):
		os.makedirs(output_path)

	light_and_texture = list(product(light_xmls, textures_path))
	num = len(light_and_texture)
	light_and_texture = np.array(light_and_texture)

	for model_path in models_path:
		count = 0
		model_name = re.match(r'.*model\\(.*).serialized', model_path)[1]
		print("Processing ", model_name)

		xml_output = output_path + "xml/{}/".format(model_name)

		if not os.path.exists(xml_output):
			os.makedirs(xml_output)

		sample_idx = np.random.choice(np.arange(num), sample_num, replace=False)
		sample_light_texture = light_and_texture[sample_idx]
		for light_xml, texture_path in sample_light_texture:
			xml_tree = ET.parse(light_xml)

			xml_tree = configure_xml(xml_tree, model_path, texture_path, scale=4)
			xml_tree = light_noise(xml_tree)
			
			xml_tree.write(xml_output + "{}.xml".format(count))
			count += 1

			if count % 100 == 0:
				print("Processed progress: ", count)
				

	print("Finished")


def batch_rendering(render_path, xml_path, output_path, render_range=None):
	if not os.path.exists(output_path):
		os.makedirs(output_path)

	t1 = time()
	xml_path = glob.glob(xml_path + "*.xml")

	if render_range is not None:
		a, b = render_range
		xml_path = xml_path[a : b]

	for path in xml_path:
		name = re.match(r'.*xml/.*\\(.*).xml', path)[1]
		os.system("{}/mitsuba.exe -x {} -o {}/{}.jpg".format(render_path, path, output_path, name))

	print("Finished, time cost: ", time() - t1)
	
	log_name = strftime("%m_%d_%H", localtime())
	with open(log_name + ".txt", 'w') as f:
		if render_range is not None:
			a, b = render_range
			print("Processed, {}, {} ~ {}".format(xml_path, a, b), file=f)
		else:
			print("Processed, {}".format(xml_path), file=f)

		print("Time cost: {} seconds".format(time()-t1), file=f)


def render_all(render_path, xml_path, output_path):
	t1 = time()
	xml_path = glob.glob(xml_path + "**/*.xml", recursive=True)

	for num, path in enumerate(xml_path):
		t2 = time()
		res = re.match(r'.*xml\\(.*)\\(.*).xml', path)
		model, name = res[1], res[2]
		img_path = output_path + model + "/"

		if not os.path.exists(img_path):
			os.makedirs(img_path)

		os.system("{}/mitsuba.exe -x -q {} -o {}/{}.jpg".format(render_path, path, img_path, name))
		print("{} finished, time cost: {}".format(path, time()-t2))

	print("Finished, time cost: ", time() - t1)
	
	log_name = strftime("%m_%d_%H", localtime())
	with open(log_name + ".txt", 'w') as f:
		print("Processed, {}".format(xml_path), file=f)

		print("Time cost: {} seconds".format(time()-t1), file=f)	


def demo(render_path):
	t1 = time()

	os.system("{}/mitsuba.exe -x ./demo/test.xml -o ./demo/test.jpg".format(render_path))

	print("Finished, time cost: ", time()-t1)


def rename():
	model_path = "./resource/model_rd/"

	models_name = glob.glob(model_path + "*.serialized")

	for num, model_name in enumerate(models_name):
		# name = re.match(r'.*\\(db_[0-9]+)_.*.serialized', model_name)[1]
		os.rename(model_name, model_path + "{}.serialized".format(num))


def extract_imgs():
	import shutil
	img_path1 = "./database/cb/"
	img_path2 = "./database/rd/"
	for num, img in enumerate(glob.glob(img_path1 + "**/*.png", recursive=True)):
		shutil.move(img, "./database/res/cb/%s.png" % num)

	for num, img in enumerate(glob.glob(img_path2 + "**/*.png", recursive=True)):
		shutil.move(img, "./database/res/rd/%s.png" % num)

if __name__ == "__main__":
	print("lalala")

	resource_path = "./resource/"
	render_path = "D:/Mitsuba"
	output_path = "./database/"

	model_name = "db_0_stripe"
	xml_path = "./database/"

	# generate_xml_sample(resource_path, output_path + "rd/", sample_num=10)
	# render_range = None
	# batch_rendering(render_path, xml_path, output_path + "images/{}/".format(model_name), render_range)
	# render_all use the follwing
	render_all(render_path, xml_path, output_path + "images/")
	# render_all(render_path, xml_path + "rd/", output_path + "rd/images/")
	# demo(render_path)